from scipy.spatial.transform import Rotation as R
from tqdm import tqdm
import json
import yaml
import tqdm
import os

with open('flatten.yaml') as f:
    flatten_conf = yaml.safe_load(f)
IMG_ROOT_PATH = flatten_conf['IMG_ROOT_PATH']
RAW_JSON_PATH = flatten_conf['RAW_JSON_PATH']
CUSTOMER_JSON_PATH = flatten_conf['CUSTOMER_JSON_PATH']
FLATTEN_JSON_ROOT_PATH = flatten_conf['FLATTEN_JSON_ROOT_PATH']
timestamp_list = flatten_conf['TIMESTAMP_LIST']
sensor_list = list(flatten_conf['SENSOR_DICT'].keys())

if not os.path.exists(FLATTEN_JSON_ROOT_PATH):
    os.mkdir(FLATTEN_JSON_ROOT_PATH)

with open(os.path.join(CUSTOMER_JSON_PATH, '1701371976.325141.json')) as f:
    changan_data = json.load(f)

samples = changan_data['sample_list']

for sample_token in tqdm.tqdm(samples, desc='Changan to Flatten...'):
    timestamp = str(round(sample_token['timestamp'] / 1000))
    timestamp = timestamp[:-6] + '.' + timestamp[-6:]
    with open(os.path.join(FLATTEN_JSON_ROOT_PATH, f'{timestamp}.json'), 'w') as f:
        with open(os.path.join(RAW_JSON_PATH, timestamp) + '.json', 'r') as r:
            raw_data = json.load(r)
            # ego orientation
            ego_pose_orientation = raw_data[0]['ego_to_world']['orientation']
            ego_pose_orientation = [ego_pose_orientation[3], ego_pose_orientation[0],
                                    ego_pose_orientation[1], ego_pose_orientation[2]]
            # ego position
            ego_pose_position = raw_data[0]['ego_to_world']['position']
            
            anno_list = sample_token['sample_annotation']
            
            obj_list = []
            visibility_dict = {}
            
            for sensor in sensor_list:
                sensor_anno = sample_token[sensor]['annotation']
                for anno in sensor_anno:
                    instance_token = anno['track_id']
                    if instance_token not in visibility_dict.keys():
                        visibility_dict[instance_token] = {}
                    visibility_dict[instance_token][sensor] = True
                
            
            for anno in anno_list:
                # type
                category = anno['type']
                # token id
                instance_token = anno['track_id']
                # position
                obj_position = anno['3D']['xyz']
                # size, convert to lwh
                obj_size_whl = anno['3D']['whl']
                obj_size = list(obj_size_whl)
                obj_size[0], obj_size[1], obj_size[2] = obj_size_whl[2], obj_size_whl[0], obj_size_whl[1]
                # rotation
                yaw = anno['3D']['yaw']
                obj_euler = [0, 0, yaw]
                # visibility
                visibility = {}
                for sensor in sensor_list:
                    if instance_token in visibility_dict.keys() and sensor in visibility_dict[instance_token].keys():
                        visibility[sensor] = True
                    else:
                        visibility[sensor] = False
                        
                obj_list.append({
                    'instance_token' : instance_token,
                    'category' : category,
                    'subcategory' : None,
                    'visibility' : visibility,
                    'obj_position' : obj_position,
                    'obj_euler' : obj_euler,
                    'obj_size' : obj_size
                })

        flatten_json_dict = {
            'ego_pose_orientation' : ego_pose_orientation,
            'ego_pose_position' : ego_pose_position,
            'objects' : obj_list
        }
        
        json.dump(flatten_json_dict, f)
    
    
    

    
        